# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2024/8/16 7:30
# File     : extract_keyword.py
# Project  : subject-word-extraction
# Desc     :

import sys
import os
sys.path.append(os.getcwd())

import pandas as pd

from scripts.cut_word.jieba_cut_word import JiebaCutWord
from scripts.keyword.IDF import IDFKeywordRxtractor
from scripts.stop_word.stop_words import load_stopwords

if __name__ == "__main__":
    jieba_cut_word = JiebaCutWord()
    df = pd.read_csv("/home/ubuntu/code/git/subject-word-extraction/data/output/2023年上市公司年度报告汇总.csv.gz")
    df['stock_code'] = df['stock_code'].astype(str).apply(lambda x: x.zfill(6))

    # df['sme_code'] = df['stock_code'].apply(lambda x: '新三板公司（全国中小企业股份转让系统）' if x[:3] == '430' or x[:3] == '830' else '')

    # df['sme_code'] = ''
    # df.loc[df['stock_code'].apply(lambda x:x[:3] in ['600','601']).index,]['sme_code'] = '主板公司（上海证券交易所）'
    # df.loc[df['stock_code'].apply(lambda x:x[:3] in ['000']).index,]['sme_code'] = '主板公司（深圳证券交易所）'
    # df.loc[df['stock_code'].apply(lambda x:x[:3] in ['002']).index,]['sme_code'] = '中小企业板（深圳证券交易所）'
    # df.loc[df['stock_code'].apply(lambda x:x[:3] in ['300']).index,]['sme_code'] = '创业板（深圳证券交易所）'

    # df['cut_words'] = df['file_content'].apply(jieba_cut_word.segment)
    
    with open("/home/ubuntu/code/git/subject-word-extraction/data/output/cut_words.txt", "w") as f:
        for i in range(df.shape[0]):
            print(f"第{i+1}行\n")
            words = jieba_cut_word.segment(df['file_content'].values[i])
            cut_words = " ".join([word for word in words if word not in load_stopwords()+['\n','\t','\r']])
            f.write(cut_words+"\n")

    # documents = df['file_content'].tolist()[:10]

    # 创建关键词提取器实例
    # extractor = IDFKeywordRxtractor()
    # # 计算 TF-IDF
    # extractor.fit(documents)
    # # 提取关键词
    # for i in range(len(documents)):
    #     print(f"文档 {i + 1} 的关键词：", extractor.extract_keywords(i))

